import os
import logging
from typing import Dict, Any

import numpy as np
import cv2
from joblib import load

import gradio as gr


logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

MODEL_PATH = os.path.join("models", "best_model.pkl")


def load_model():
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"未找到模型文件：{MODEL_PATH}，请先运行 train_ensemble.py 进行训练保存。")
    model = load(MODEL_PATH)
    logging.info(f"已加载模型：{MODEL_PATH}")
    return model


MODEL = load_model()


LABELS = {0: "猫", 1: "狗"}


def preprocess_to_flat(img_np: np.ndarray) -> np.ndarray:
    """将 RGB numpy 图像转为 32x32 灰度并展平。"""
    if img_np is None:
        raise ValueError("未接收到图像")
    if img_np.ndim == 2:  # 灰度
        gray = img_np
    else:
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    gray = cv2.resize(gray, (32, 32))
    feat = gray.reshape(1, -1).astype("float32")
    return feat


def predict(img_np: np.ndarray) -> Dict[str, Any]:
    x = preprocess_to_flat(img_np)
    pred = int(MODEL.predict(x)[0])
    label = LABELS.get(pred, str(pred))

    conf = None
    if hasattr(MODEL, "predict_proba"):
        try:
            prob = float(MODEL.predict_proba(x)[0][pred])
            conf = prob
        except Exception:
            conf = None

    if conf is None:
        return {"result": label}
    else:
        return {"result": f"{label} ({conf:.2%})"}


with gr.Blocks(title="Cats vs Dogs - Ensemble") as demo:
    gr.Markdown("# 猫狗分类（集成模型）\n上传一张图片，模型将预测猫还是狗。仅 CPU，flat 特征预处理。")
    with gr.Row():
        inp = gr.Image(type="numpy", label="上传图片")
        out = gr.Textbox(label="预测结果")
    btn = gr.Button("预测")
    btn.click(predict, inputs=inp, outputs=out)


if __name__ == "__main__":
    # 默认本地运行
    demo.launch()
